package com.qiongqi.datadir.helper;

import com.bettercloud.vault.Vault;
import com.bettercloud.vault.VaultException;
import com.bettercloud.vault.response.LogicalResponse;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.stereotype.Component;
import paillierp.Paillier;
import paillierp.key.PaillierPrivateThresholdKey;

import java.io.*;
import java.util.*;

import static java.util.Objects.nonNull;
import static org.apache.commons.lang3.StringUtils.isNotBlank;

@Slf4j
@Component
public final class VaultHelper {

    @Autowired
    @Qualifier("vaultClient")
    private Vault vaultClient;

    /**
     * 存储私钥分片（仅用于初始化，不应频繁调用）
     */
    public boolean storePrivateKeyShare(String bucketName, String shareName, PaillierPrivateThresholdKey share) throws VaultException {
        String path = "secret/data/" + bucketName + "/shares/" + shareName;
        try(ByteArrayOutputStream bos = new ByteArrayOutputStream(); ObjectOutputStream oos = new ObjectOutputStream(bos)){
            oos.writeObject(share);
            String serialized = Base64.getEncoder().encodeToString(bos.toByteArray());
            final Map<String, Object> secrets = new HashMap<>();
            secrets.put("value", serialized);
            final LogicalResponse writeResponse = vaultClient.logical()
                    .write(path, secrets);
            return nonNull(writeResponse) && isNotBlank(writeResponse.getLeaseId());
        }catch(IOException ex){
            log.error("Failed to serialize private key:{}", ex);
        }
        return false;
    }

    /**
     * 存储分片密钥中的公钥（仅用于初始化，不应频繁调用）
     */
    public boolean storePublicKeyShare(String bucketName, Paillier publicKey) throws VaultException {
        String path = "secret/data/" + bucketName + "/pubkey";
        try(ByteArrayOutputStream bos = new ByteArrayOutputStream(); ObjectOutputStream oos = new ObjectOutputStream(bos)){
            oos.writeObject(publicKey);
            String serialized = Base64.getEncoder().encodeToString(bos.toByteArray());
            final Map<String, Object> secrets = new HashMap<>();
            secrets.put("value", serialized);
            //log.info("store path:{}，serialized:{} ", path, serialized);
            final LogicalResponse writeResponse = vaultClient.logical()
                    .write(path, secrets);
            return nonNull(writeResponse) && isNotBlank(writeResponse.getLeaseId());
        }catch(IOException ex){
            log.error("Failed to serialize private key:{}", ex);
        }
        return false;
    }

    /**
     * 读取私钥分片（仅用于本地计算，不应对外暴露）
     */
    public PaillierPrivateThresholdKey getPrivateKeyShare(String bucketName, String shareName) throws VaultException {
        String path = "secret/data/" + bucketName + "/shares/" + shareName;
        final String serialized = vaultClient.logical()
                .read(path)
                .getData().get("value");
        byte[] bytes = Base64.getDecoder().decode(serialized);
        try(ObjectInputStream ois = new ObjectInputStream(new ByteArrayInputStream(bytes))){
            return (PaillierPrivateThresholdKey) ois.readObject();
        }catch(Exception ex){
            log.error("Failed to deserialize private key:{}", ex);
        }
        return null;
    }

    /**
     * 读取私钥分片（仅用于本地计算，不应对外暴露）
     */
    public Paillier getPublicKeyShare(String bucketName) throws VaultException {
        String path = "secret/data/" + bucketName + "/pubkey";
        final String serialized = vaultClient.logical()
                .read(path)
                .getData().get("value");
        //log.info("get path:{}，serialized:{} ", path, serialized);
        byte[] bytes = Base64.getDecoder().decode(serialized);
        try(ObjectInputStream ois = new ObjectInputStream(new ByteArrayInputStream(bytes))){
            return (Paillier) ois.readObject();
        }catch(Exception ex){
            log.error("Failed to deserialize private key:{}", ex);
        }
        return null;
    }
}

